import cv2
import math
import numpy as np

#以rgb格式读取图片文件，opencv默认读取的图片为BGR
def read_image_rgb(imgPath):
    img = cv2.imread(imgPath)
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    return img

#以rgb格式读取图片文件,channel放在前面，（channel,H,W）
def read_image_channel_first(imgPath):
    img = read_image_rgb(imgPath)
    img = img.transpose(2,0,1)
    return img

#显示rgb图片
def show_rgb(rgb_img):
    img = cv2.cvtColor(rgb_img,cv2.COLOR_RGB2BGR)
    cv2.imshow('rgb image',img)
    cv2.waitKey()

#将rgb图片转成bgr图片，方便使用opencv保存图片
def rgb2bgr(rgb_img):
    img = cv2.cvtColor(rgb_img,cv2.COLOR_RGB2BGR)
    return img

#将形状为(channel,H,W)，图片转化为(H,W，channel)
def channel_first_to_last(img):
    img = img.transpose(1,2,0)
    return img

#将形状为(channel,H,W)，图片转化为(H,W，channel)
def channel_last_to_first(img):
    img = img.transpose(2,0,1)
    return img

#bicubic缩放图片
def resize_bicubic(img,h,w):
    img = cv2.resize(img,(w,h),interpolation=cv2.INTER_CUBIC)
    return img

#切割图片成指定大小的小块,可用于BGR或RGB图片上
def seg_img(img,h,w):
    size  = img.shape[:2]
    imgs = []
    a = size[0]//h
    b = size[1]//w
    for y in range(a):
        for x in range(b):
            seg = img[y*h:(y+1)*h,x*w:(x+1)*w]
            imgs.append(seg)
    return imgs

#保存rgb图片
def save_rgb(rgb_img,savePath):
    img = rgb2bgr(rgb_img)
    cv2.imwrite(savePath,img)
    
#参考：https://blog.csdn.net/qazwsxrx/article/details/104550550
def psnr(img1, img2):
    img1 = np.float64(img1)
    img2 = np.float64(img2)
    mse = np.mean( (img1 - img2) ** 2 )
    if mse == 0:
        return 100
    PIXEL_MAX = 255.0
    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))

#参考：https://blog.csdn.net/a2824256/article/details/115013851
def ssim(img1, img2):
    """Calculate SSIM (structural similarity) for one channel images.
    It is called by func:`calculate_ssim`.
    Args:
        img1 (ndarray): Images with range [0, 255] with order 'HWC'.
        img2 (ndarray): Images with range [0, 255] with order 'HWC'.
    Returns:
        float: ssim result.
    """

    C1 = (0.01 * 255)**2
    C2 = (0.03 * 255)**2

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
	# 公式二计算
    ssim_map = ((2 * mu1_mu2 + C1) *
                (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
                                       (sigma1_sq + sigma2_sq + C2))
    return ssim_map.mean()

#cv2.flip使用 参考: https://blog.csdn.net/nienelong3319/article/details/95872454
#图片翻转，返回(原图，翻转90度，水平翻转的图片)
def flip(img):
    rotation90 = np.rot90(img) #翻转90度
    horizontal = cv2.flip(img,1) #图像水平翻转
    return (img,rotation90,horizontal)


#progress bar
def print_progress_bar(now,end,barsize=50):
    p = now/end
    bar = "\r["
    bar +="-"*int((barsize*p))
    bar +=">"
    bar +=" "*int(barsize-len(bar))
    bar +="]"
    bar += str(int(p*100))
    bar += "%"
    print(bar,end="")
    

if __name__ == "__main__":
    img = read_image_rgb('E:/Datasets/anime HR/f656ed39697f7c61742df35ef1bd7e26.jpg')
    #img = resize_bicubic(img, w=800, h=700)
    segs = seg_img(img,192,192)
    fImg = flip(segs[9])
    for i in range(len(fImg)):
        cv2.imwrite("test/%d.jpg"%(i),fImg[i])
    
    
    




